import numpy as np
import matplotlib.pyplot as plt
import math

# 定义一个函数，用于生成正弦波
def getSin(amp, freq, phase, sampleList):
    return amp * np.sin(-2 * math.pi * freq * sampleList + phase)

# 定义一个函数，用于生成余弦波
def getCos(amp, freq, phase, sampleList):
    return amp * np.cos(-2 * math.pi * freq * sampleList + phase)

# 定义去噪函数，将振幅小于阈值的噪声去除
def denoise(arr, thresh):
    mask = arr > thresh
    return mask * arr

# 设置采样率为 3000
srate = 3000
# 在 0 到 1 的范围内生成等间距的采样点
t = np.linspace(0, 1, srate)

# 生成四个正弦和余弦波
signals = []
signals.append(getSin(amp=1.5, freq=30, phase=0, sampleList=t))
signals.append(getCos(amp=3, freq=5, phase=0, sampleList=t))
signals.append(getSin(amp=10, freq=100, phase=0, sampleList=t))
signals.append(getCos(amp=20, freq=120, phase=0, sampleList=t))

# 将四个正弦和余弦波相加，得到混合信号
mixed_signal = np.zeros_like(signals[0])
for signal in signals:
    mixed_signal += signal

# 使用 numpy 的 fft 函数计算傅里叶系数
fCoefs = np.fft.fft(mixed_signal)

# 计算振幅列表
amp_list = np.abs(fCoefs) / len(mixed_signal) * 2

# 将振幅数组和频率数组进行平移，使 0 Hz 位于中心
freqs = np.fft.fftfreq(len(amp_list), 1 / srate)
amp_shifted = np.fft.fftshift(amp_list)
freq_shift = np.fft.fftshift(freqs)

# 创建两个子图，分别用于显示原始振幅频谱和处理后的振幅频谱
fig, ax = plt.subplots(1, 2, figsize=(12, 4))

# 绘制原始振幅频谱
ax[0].stem(freq_shift, amp_shifted)
ax[0].set_xlim([-150, 150])
ax[0].set_ylim([-0.5, 21])
ax[0].grid()

# 对振幅频谱进行去噪处理
for i in range(len(freq_shift)):
    if freq_shift[i] > 110 or freq_shift[i] < -110:
        amp_shifted[i] = 0
amp_shifted = denoise(amp_shifted, 1)

# 绘制处理后的振幅频谱
ax[1].stem(freq_shift, amp_shifted)
# 设置第二个子图的坐标轴范围
ax[1].set_xlim([-150, 150])
ax[1].set_ylim([0, 20])
ax[1].set_yticks(np.arange(0, 21, 2))
ax[1].grid()
#在第二个子图中标记去除的频率点
for i in range(len(freq_shift)):
      if freq_shift[i] > 110 or freq_shift[i] < -110:
         ax[1].annotate('X', xy=(freq_shift[i], amp_shifted[i]), ha='center', va='center', color='r', fontsize=12)

#显示图像
plt.show()
